import numpy as np
import matplotlib.pyplot as plt
import pywt
from scipy.signal import stft, hilbert
from scipy.fftpack import fft

# 生成模拟数据
# 生成模拟数据
fs = 1000  # 采样频率
t = np.linspace(0, 1, fs, endpoint=False)
amplitude = 1.5 # 基础振幅
np.random.seed(0)  # 确保每次运行代码生成相同的随机数

# 创建更复杂的信号
signal = sum(amplitude * np.random.rand() * np.sin(2 * np.pi * np.random.randint(30, 100) * t + np.random.rand())
             for _ in range(5))  # 随机生成5个正弦波的和

noise = np.random.normal(0, 0.2, fs)  # 高斯噪声

# 生成白噪声
noise_amplitude = 0.5  # 噪声的振幅
white_noise = noise_amplitude * np.random.normal(size=t.shape)  # 均值为0，标准差为0.5的正态分布
mixed_signal = signal + noise  # 含噪声的信号

# 第一阶段处理：小波变换去噪和FFT

coeffs = pywt.wavedec(mixed_signal, 'db4', level=4)
threshold = 0.1
thresholded_coeffs = [pywt.threshold(i, threshold, mode='soft') for i in coeffs]
wavelet_signal = pywt.waverec(thresholded_coeffs, 'db4')
fft_signal = fft(wavelet_signal)   # 计算FFT

# 第二阶段处理：SFFT和希尔伯特变换
f, t_, Zxx = stft(wavelet_signal, fs, nperseg=100)
hilbert_signal = hilbert(wavelet_signal)
envelope_signal = np.abs(hilbert_signal)








# 计算信噪比
def calculate_snr(signal, noise):
    signal_power = np.mean(signal**2)
    noise_power = np.mean(noise**2)
    return 10 * np.log10(signal_power / noise_power)

snr_wavelet = calculate_snr(wavelet_signal, mixed_signal - wavelet_signal)  # 计算信噪比
snr_fft = calculate_snr(fft_signal, wavelet_signal - fft_signal)  # 计算信噪比
snr_envelope = calculate_snr(envelope_signal, mixed_signal - envelope_signal)  # 计算信噪比


#信号功率
signal_power = np.mean(signal**2)
#噪声功率
noise_power = np.mean(noise**2)
#信噪比
snr = 10 * np.log10(signal_power / noise_power)

# 输出信号功率和噪声功率
print(f"Signal Power: {signal_power:.2f}")
print(f"Noise Power: {noise_power:.2f}")

# 输出信噪比结果
print(f"SNR of original signal: {calculate_snr(signal, noise):.2f} dB")
print(f"SNR after wavelet processing: {snr_wavelet:.2f} dB")   # 输出信噪比结果

print(f"SNR after FFT processing: {snr_fft:.2f} dB")  # 输出信噪比结果
print(f"SNR after envelope extraction: {snr_envelope:.2f} dB")  # 输出信噪比结果

# 绘图
plt.figure(figsize=(14, 10))
plt.subplot(411)
plt.plot(t, mixed_signal, label="Original Signal")
plt.title("Original Signal")
plt.subplot(412)
plt.plot(t, wavelet_signal, label="Wavelet Processed Signal")
plt.title("Wavelet Processed Signal")
plt.subplot(413)
plt.pcolormesh(t_, f, np.abs(Zxx), shading='gouraud')
plt.title("SFFT of Wavelet Processed Signal")
plt.ylabel("Frequency [Hz]")
plt.subplot(414)

plt.title("Envelope Extracted using Hilbert Transform")
plt.tight_layout()
plt.show()
